#include <linux/types.h>
#include <linux/kernel.h>
#include <linux/netfilter.h>
#include <linux/module.h>
#include <linux/skbuff.h>
#include <linux/ip.h>
#include <linux/in.h>
#include <linux/proc_fs.h>
#include <linux/vmalloc.h>
#include <linux/err.h>
#include <net/netfilter/nf_conntrack.h>
#include <linux/netfilter_ipv4.h>

#include <net/sock.h>

#include "ip_notify.h"
#include "ip_notify_pub.h"


MODULE_LICENSE( "GPL" );
MODULE_AUTHOR( "Vinson Yan <vinson.one.yan@gmail.com>" );

DEFINE_SPINLOCK( ip_notify_lock );
DECLARE_WAIT_QUEUE_HEAD(ip_notify_wait);
static atomic_t notify_num = ATOMIC_INIT(0);

struct hlist_head *notify_hash __read_mostly;
static int ipnotify_vmalloc  __read_mostly;
static struct notify_to_user * g_ip_notify;

static void ip_notify_free( struct ip_notify *notify )
{
	if ( !( atomic_dec_and_test( &notify->refcnt) ) )
		kfree( notify );
}

static struct hlist_head *alloc_hash_table(int size, int *vmalloced)
{
	struct hlist_head *hash;
	unsigned int i;

	*vmalloced = 0;
	hash = (void*)__get_free_pages(GFP_KERNEL,
			get_order(sizeof(struct hlist_head) * size));
	if (!hash) {
		*vmalloced = 1;
		hash = vmalloc(sizeof(struct hlist_head) * size);
	}

	if (hash)
		for (i = 0; i < size; i++)
			INIT_HLIST_HEAD(&hash[i]);

	return hash;
}

static void free_hash_table( struct hlist_head *hash, int vmalloced, int size )
{
	int i;
	struct ip_notify     *curr;
	struct hlist_node *h, *n;    

	spin_lock_bh( &ip_notify_lock );
	for ( i = 0; i < size; i++ ) {
		hlist_for_each_entry_safe(curr, h, n, &hash[i], hnode) {
			hlist_del( &curr->hnode );
			ip_notify_free( curr );
		}
	}
	spin_unlock_bh( &ip_notify_lock );

	if (vmalloced)
		vfree(hash);
	else
		free_pages((unsigned long)hash,
				get_order(sizeof(struct list_head) * size));
}

#if 0
static void ip_notify_wake_up( wait_queue_head_t wait )
{	
	wake_up_interruptible( &wait );
}
#endif

static u32 ip_notify_hash( u32 sip )
{
	return ntohl( sip ) & 0xff;
}

static u32 ip_notify_new( u32 sip, u32 dip, u16 sport, u16 dport, u8 protocol )
{
	u32 hash = ip_notify_hash( sip );
	struct ip_notify *node = ( struct ip_notify * )kzalloc( sizeof( struct ip_notify ), GFP_ATOMIC );

	if ( !node )
		return -1;

	node->sip = sip ;
	node->dip = dip ;
	node->sport = sport;
	node->dport = dport;
	node->protocol = protocol;
	INIT_HLIST_NODE( &node->hnode );

	spin_lock_bh( &ip_notify_lock );
	hlist_add_head( &node->hnode, &notify_hash[hash] );
	spin_unlock_bh( &ip_notify_lock );

	atomic_set( &node->refcnt, 1 );

	printk( "1 cur->sip [ 0x%x ], dip [ 0x%x ], sport [ %u ], dport [ %u ].\n", 
		ntohl( node->sip ), ntohl( node->dip ), ntohs( node->sport ), ntohs( node->dport ) );

	memcpy( g_ip_notify, node, sizeof( struct ip_notify ) );
	
	printk( "2 cur->sip [ 0x%x ], dip [ 0x%x ], sport [ %u ], dport [ %u ].\n", 
			ntohl( g_ip_notify->sip ), ntohl( g_ip_notify->dip ), 
			ntohs( g_ip_notify->sport ), ntohs( g_ip_notify->dport ) );

	atomic_inc( &notify_num );
	wake_up_interruptible( &ip_notify_wait );

	return 0;
}

static u32 ip_notify_create( struct sk_buff * skb, struct nf_conn *ct, u8 dir )
{
	u32     sip = 0;
	u32     dip = 0;
	u16     sport = 0;
	u16     dport = 0;
	u8      protocol     = ip_hdr( skb )->protocol;
	struct hlist_node *node;
	struct ip_notify *curr = NULL;
	u32     hash = 0;

	sip = ct->tuplehash[dir].tuple.src.u3.ip;
	dip = ct->tuplehash[dir].tuple.dst.u3.ip;

	if ( protocol == IPPROTO_TCP ) {
		sport = ct->tuplehash[dir].tuple.src.u.tcp.port;
		dport = ct->tuplehash[dir].tuple.dst.u.tcp.port;
	} else if ( protocol == IPPROTO_UDP ) {
		sport = ct->tuplehash[dir].tuple.src.u.udp.port;
		dport = ct->tuplehash[dir].tuple.dst.u.udp.port;         
	} else
		return 0;

	printk( "sip 0x%x, dip 0x%x, sport %u, dport %u.\n", sip, dip, sport, dport );

	hash = ip_notify_hash( sip );

	spin_lock_bh( &ip_notify_lock );
	hlist_for_each_entry( curr, node, &notify_hash[hash], hnode ) {
		if ( curr->sip == sip && curr->dip == dip && curr->sport == sport && curr->dport == dport ) {
			printk( "Find OK.\n" );
			spin_unlock_bh( &ip_notify_lock );
			return 0;
		}
	}
	spin_unlock_bh( &ip_notify_lock );

	printk( "sip 0x%x, dip 0x%x, sport %u, dport %u.\n", sip, dip, sport, dport );
	ip_notify_new( sip, dip, sport, dport, protocol );

	return 0;
}

static unsigned int ipv4_notify_hook( unsigned int hooknum,
		struct sk_buff *skb,
		const struct net_device *in,
		const struct net_device *out,
		int ( *okfn )( struct sk_buff * ) )
{
	u32     target = 1;
	u8     dir = 0;
	//struct ip_notify *notify = NULL;
	struct nf_conn     *ct;
	enum ip_conntrack_info     ctinfo;
	u8 protocol = ip_hdr( skb )->protocol;

	printk( "%s %d protocol[%u].\n", __FUNCTION__, __LINE__, protocol );

	if ( IPPROTO_TCP != protocol && IPPROTO_UDP != protocol )
		return target;

	ct = nf_ct_get( skb, &ctinfo );
	if ( !ct )
		return target;

	dir = CTINFO2DIR( ctinfo );

	ip_notify_create( skb, ct, dir );

	return target;

}

static struct nf_hook_ops ipv4_notify_ops[] = {
	{
		.hook          = ipv4_notify_hook,
		.owner          = THIS_MODULE,
		.pf          = PF_INET,
		.hooknum     = NF_IP_FORWARD,
		.priority     = NF_IP_PRI_CONNTRACK + 1,
	},
};

static int ip_notify_open( struct inode *inode, struct file *filp )
{
	struct ip_notify_file     *status;

	status = kzalloc( sizeof( struct ip_notify_file ),  GFP_ATOMIC );
	if ( NULL == status )
		return -ENOMEM;

	filp->private_data = status;
	atomic_set( &status->flags, 0 );

	return 0;
}

static int ip_notify_read( struct file *file, char __user *buf,
		size_t len, loff_t *ppos) 
{
	struct ip_notify_file *status = file->private_data;
	u32	ret;

	if ( !status )
		return -ENXIO;
#if 0
	DECLARE_WAITQUEUE( wait,  current );

	add_wait_queue( &ip_notify_wait, &wait );
	while ( 1 ) {
		set_current_state( TASK_INTERRUPTIBLE );
		atomic_set( &status->flags, 1 );

		if ( atomic_read( &notify_num ) )
			break;
		ret = -EAGAIN;

		if ( file->f_flags & O_NONBLOCK )
			break;

		if ( signal_pending( current ))
			break;
		schedule();
	}

	set_current_state( TASK_RUNNING );
	remove_wait_queue( &ip_notify_wait, &wait );
#endif
	if ( file->f_flags & O_NONBLOCK )
		return -1;

	wait_event_interruptible( ip_notify_wait, atomic_read( &notify_num ) );

	atomic_set( &notify_num, 0 );

	printk( " notify_num %u, len  %u.\n", atomic_read( &notify_num ), len );
	ret = -EOVERFLOW;

	if ( len < sizeof( struct notify_to_user ) )
		goto out;

	if ( copy_to_user( buf,  g_ip_notify, len ) )
		goto outf;


	ret = len;
out:
	atomic_set( &status->flags, 0 );

outf:
	memset( g_ip_notify, 0, sizeof( struct ip_notify ) );

	return ret;
}

static int ip_notify_release( struct inode *inode, struct file *file )
{
	struct ip_notify_file *status = file->private_data;

	if ( status ) {
		file->private_data = NULL;
		kfree( status );
	}

	return 0;
}

static const struct file_operations ip_notify_fops = {
	.owner               = THIS_MODULE,
	.open          = ip_notify_open,
	.read          = ip_notify_read,
	.release             = ip_notify_release,
};

static int __init ip_notify_init( void )
{
	u32     retu = 0;

	nf_register_hooks( ipv4_notify_ops, ARRAY_SIZE( ipv4_notify_ops ) );

	if( !proc_net_fops_create( &init_net, "ipnotify", S_IRUGO, &ip_notify_fops ) ) {
		retu = -1;
		goto un_register_hook;
	}

	notify_hash = alloc_hash_table( 256, &ipnotify_vmalloc );
	if ( NULL == notify_hash ) {
		retu = -1;
		goto proc_net_fops_destory;
	}

	g_ip_notify = kzalloc( sizeof( struct notify_to_user ), GFP_ATOMIC );
	if ( NULL == g_ip_notify ) {
		retu = -1;
		goto free_hash;
	}
	printk( "Initiation notify module success.\n" );

	return retu;

free_hash:
	free_hash_table( notify_hash, ipnotify_vmalloc, 256 );

proc_net_fops_destory:
	proc_net_remove(&init_net, "ipnotify");

un_register_hook:
	nf_unregister_hooks( ipv4_notify_ops, ARRAY_SIZE( ipv4_notify_ops ) );

	return retu;

}

static void __exit ip_notify_exit( void )
{
	nf_unregister_hooks( ipv4_notify_ops, ARRAY_SIZE( ipv4_notify_ops ) );
	proc_net_remove(&init_net, "ipnotify");
	free_hash_table( notify_hash, ipnotify_vmalloc, 256 );
	kfree( g_ip_notify );
	printk( "Exit notify module success.\n" );

	return ;
}

module_init( ip_notify_init );
module_exit( ip_notify_exit );
